import pandas as pd
import numpy as np
from sklearn.model_selection import StratifiedKFold
from sklearn import metrics
from sklearn.ensemble import GradientBoostingClassifier
import sys
import warnings
warnings.filterwarnings('ignore')

params = {
    'n_estimators': 5000,
    'max_depth': 100,
}
model = GradientBoostingClassifier(**params)

if len(sys.argv) > 1:
    n = int(sys.argv[1])
else:
    n = 100
trn_file = f'trn_per_{n}.csv'
print(trn_file)

trn_data = pd.read_csv(f'./data/{trn_file}', header=0)
X = trn_data.iloc[:, :-2]
y = trn_data['type']
X1 = pd.read_csv('./data/tst_desc.csv', header=0)

A = pd.read_csv('fts/fea.csv', header=0)

nfea_arr = list(range(40, 51))
oof_arr = []
for nn in nfea_arr:
    Nfea, nfea = len(A), nn
    features = A['fea_name'][Nfea-nfea:].tolist()
    print(len(features))

    fold = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    models = []
    pred = np.zeros((len(X1), 3))
    oof = np.zeros((len(X), 3))
    for index, (train_idx, val_idx) in enumerate(fold.split(X, y)):
        Xt, yt = X.loc[train_idx, features], y.iloc[train_idx]
        Xv, yv = X.loc[val_idx, features], y.iloc[val_idx]
        model = model.fit(Xt, yt)
        models.append(model)
        val_pred = model.predict_proba(Xv)
        oof[val_idx] = val_pred
        val_y = y.iloc[val_idx]
        val_pred = np.argmax(val_pred, axis=1)
        print(index, 'val f1', metrics.f1_score(val_y, val_pred, average='macro'))

    oof = np.argmax(oof, axis=1)
    all_f1 = metrics.f1_score(oof, y, average='macro')
    print('oof f1', all_f1)
    oof_arr.append(all_f1)

pd.DataFrame(np.array([nfea_arr, oof_arr]), index=['num', 'f1_acc']).T.to_csv('fts/f1acc40_50.csv', index=False)


import matplotlib.pyplot as plt
plt.plot(nfea_arr, oof_arr)
plt.grid()
plt.savefig('fts/nfea_acc40_50.png')
